#!/usr/bin/env python3
# Copyright (c) 2024 maminjie <canpool@163.com>
# SPDX-License-Identifier: MulanPSL-2.0

"""openEuler Security Bulletins"""

import urllib
import urllib.request
import urllib.parse
import urllib.error
import json
import csv
import argparse
import datetime
import os


HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/125.0.0.0 Safari/537.36 Edg/125.0.0.0"
}
SECURITY_URL = "https://www.openeuler.org/api-euler/api-cve/cve-security-notice-server/securitynotice"
FINDALL_URL = SECURITY_URL + "/findAll"
DETAIL_URL = SECURITY_URL + "/getBySecurityNoticeNo?securityNoticeNo={}"
CVE_DB_URL = "https://www.openeuler.org/api-euler/api-cve/cve-security-notice-server/cvedatabase"
CVE_URL = CVE_DB_URL + "/getByCveIdAndPackageName?cveId={}&packageName={}"


def url_get(url):
    req = urllib.request.Request(url=url, headers=HEADERS)
    try:
        result = urllib.request.urlopen(req)
        return result.read().decode("utf-8")
    except (urllib.error.HTTPError, urllib.error.URLError) as e:
        print(e.reason)
        return None


def url_post(url, values):
    data = values.encode('utf-8')
    req = urllib.request.Request(url=url, data=data, headers=HEADERS, method="POST")
    try:
        result = urllib.request.urlopen(req)
        return result.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        print(str(e))
        return None


def find_all(page, size):
    req_data = {
        "pages": {
            "page": page,
            "size": size
        },
        "keyword": "",
        "type": [],
        "date": [],
        "affectedProduct": [],
        "affectedComponent": "",
        "noticeType": "cve"
    }
    resp_str = url_post(FINDALL_URL, json.dumps(req_data))
    if not resp_str:
        return None
    resp_data = json.loads(resp_str)
    if resp_data.get("code", 1) != 0:
        return None
    return resp_data.get("result", None)


def get_detail(item, product, number):
    detail_str = url_get(DETAIL_URL.format(number))
    if not detail_str:
        return
    package_helper_lst = json.loads(detail_str)["result"]["packageHelperList"]
    for package_helper in package_helper_lst:
        affectedProduct = package_helper["productName"]
        if product != affectedProduct:
            continue
        child_lst = package_helper["child"]
        for child in child_lst:
            system_arch = child["productName"]
            item[system_arch] = []
            rpm_lst = child["child"]
            for rpm in rpm_lst:
                rpm_pack_name = rpm["packageName"]
                item[system_arch].append(rpm_pack_name)


def get_cve_cvss(item, cveid, package):
    item["cvss"] = ""
    cve_ids = cveid.rstrip(";").strip().split(";")
    for cve_id in cve_ids:
        resp = url_get(CVE_URL.format(cve_id, package))
        if not resp:
            continue
        cvss = json.loads(resp)["result"]["cvsssCoreOE"]
        item["cvss"] += cvss + ";"


def get_total_count():
    resp = find_all(1, 1)
    if not resp:
        return 0
    return resp.get("totalCount", 0)


def check_time(update_time, start_time="", end_time=""):
    date_format = "%Y-%m-%d %H:%M:%S"
    update_time_obj = datetime.datetime.strptime(update_time, date_format)
    if start_time:
        try:
            start_time_obj = datetime.datetime.strptime(start_time, date_format)
            if update_time_obj < start_time_obj:
                return False
        except ValueError:
            pass
    if end_time:
        try:
            end_time_obj = datetime.datetime.strptime(end_time, date_format)
            if update_time_obj > end_time_obj:
                return False
        except ValueError:
            pass
    return True


def csv_line_wrap(field):
    return "{}".format(field.replace(';', '\n'))


def do_notice(args):
    total_count = get_total_count()
    if total_count == 0:
        return 1
    if args.outfile:
        out_file = args.outfile
    else:
        out_file = "openeuler_security_notice-{}.csv".format(datetime.datetime.now().strftime("%Y%m%d"))
        out_file = os.path.join(os.getcwd(), out_file)
    csv_file_handler = open(out_file, 'w', encoding='utf-8', newline="")
    csv_writer = csv.writer(csv_file_handler)
    csv_writer.writerow(["securityNoticeNo", "summary", "type", "affectedProduct", "affectedComponent",
                         "updateTime", 'cveId', "cvss", "src.rpm", "noarch.rpm", "aarch64.rpm", "x86_64.rpm"
                         ])
    per_page_size = 50
    for i in range(round(total_count / per_page_size)):
        resp = find_all(i + 1, per_page_size)
        security_notice_lst = resp["securityNoticeList"]
        for security_notice in security_notice_lst:
            if args.product not in security_notice["affectedProduct"]:
                continue
            if not check_time(security_notice["updateTime"], args.starttime, args.endtime):
                continue
            item = {
                "securityNoticeNo": security_notice["securityNoticeNo"],
                "summary": security_notice["summary"],
                "type": security_notice["type"],
                "affectedProduct": security_notice["affectedProduct"],
                "affectedComponent": security_notice["affectedComponent"],
                "updateTime": security_notice["updateTime"],
                'cveId': security_notice['cveId']
            }
            get_detail(item, args.product, security_notice["securityNoticeNo"])
            get_cve_cvss(item, security_notice['cveId'], security_notice["affectedComponent"])
            csv_writer.writerow(
                [item["securityNoticeNo"],
                 item["summary"],
                 item["type"],
                 csv_line_wrap(item["affectedProduct"]),
                 item["affectedComponent"],
                 item["updateTime"],
                 csv_line_wrap(item['cveId'].rstrip(';')),
                 csv_line_wrap(item['cvss'].rstrip(';')),
                 csv_line_wrap(";".join(item.get("src", []))),
                 csv_line_wrap(";".join(item.get("noarch", []))),
                 csv_line_wrap(";".join(item.get("aarch64", []))),
                 csv_line_wrap(";".join(item.get("x86_64", []))), ]
            )
            print(item)
    csv_file_handler.close()
    return 0


def do_package(args):
    total_count = get_total_count()
    if total_count == 0:
        return 1
    if args.outfile:
        out_file = args.outfile
    else:
        out_file = "openeuler_security_package-{}.csv".format(datetime.datetime.now().strftime("%Y%m%d"))
        out_file = os.path.join(os.getcwd(), out_file)
    csv_file_handler = open(out_file, 'w', encoding='utf-8', newline="")
    csv_writer = csv.writer(csv_file_handler)
    csv_writer.writerow(["affectedComponent", "summary", "updateTime", 'cveId', "securityNoticeNo"])
    per_page_size = 50
    result = {}
    for i in range(round(total_count / per_page_size)):
        resp = find_all(i + 1, per_page_size)
        security_notice_lst = resp["securityNoticeList"]
        for security_notice in security_notice_lst:
            if args.product not in security_notice["affectedProduct"]:
                continue
            if not check_time(security_notice["updateTime"], args.starttime, args.endtime):
                continue
            pkg = security_notice["affectedComponent"]
            if pkg in result:
                if security_notice["updateTime"] not in result[pkg]["updateTime"]:
                    result[pkg]["updateTime"] += ";" + security_notice["updateTime"]
                result[pkg]["cveId"] += security_notice['cveId']
                result[pkg]["securityNoticeNo"] += ";" + security_notice["securityNoticeNo"]
            else:
                result[pkg] = {}
                result[pkg]["summary"] = security_notice["summary"]
                result[pkg]["updateTime"] = security_notice["updateTime"]
                result[pkg]["cveId"] = security_notice['cveId']
                result[pkg]["securityNoticeNo"] = security_notice["securityNoticeNo"]
            print(security_notice)
    pkgs = sorted(result.keys())
    for pkg in pkgs:
        csv_writer.writerow([pkg,
                             result[pkg]["summary"],
                             csv_line_wrap(result[pkg]["updateTime"]),
                             csv_line_wrap(result[pkg]['cveId'].rstrip(';')),
                             csv_line_wrap(result[pkg]["securityNoticeNo"]),
                             ])
    csv_file_handler.close()
    return 0


def do_main(args):
    print("try -h/--help for more details.")


def main():
    parser = argparse.ArgumentParser(description="openEuler Security Bulletins Tool")
    parser.set_defaults(func=do_main)

    subparsers = parser.add_subparsers(help="oesb sub-commands")
    # notice
    subparser = subparsers.add_parser("notice", aliases=["n"], help="get openEuler security notice")
    subparser.add_argument("-p", "--product", type=str, default="openEuler-22.03-LTS-SP1",
                        help="The affected openEuler product name")
    subparser.add_argument("-s", "--starttime", type=str, default="",
                        help="The start time, format is: 2024-01-31 14:31:42")
    subparser.add_argument("-e", "--endtime", type=str, default="",
                        help="The end time, format is: 2024-06-12 18:31:42")
    subparser.add_argument("-o", "--outfile", type=str, default="",
                        help="The output file, default is openeuler_security_notice-YYYYMMDD.csv in cwd")
    subparser.set_defaults(func=do_notice)
    # package
    subparser = subparsers.add_parser("package", aliases=["p"], help="get openEuler security affected component")
    subparser.add_argument("-p", "--product", type=str, default="openEuler-22.03-LTS-SP1",
                           help="The affected openEuler product name")
    subparser.add_argument("-s", "--starttime", type=str, default="",
                           help="The start time, format is: 2024-01-31 14:31:42")
    subparser.add_argument("-e", "--endtime", type=str, default="",
                           help="The end time, format is: 2024-06-12 18:31:42")
    subparser.add_argument("-o", "--outfile", type=str, default="",
                           help="The output file, default is openeuler_security_package-YYYYMMDD.csv in cwd")
    subparser.set_defaults(func=do_package)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
